/**********************************************
Hari Sundar - 8 Nov 2005
Myocyte Tracking Program

--* Need to add anyoption     *--
--* Need to handle exceptions *--
*********************************************/

#include <iostream>
#include <fstream>
#include <algorithm>
#include <vector>
#include <string>
#include<ctime>

// #include <omp.h>

#include "Point.h"

class triplet {
public:
	int i;
	int j;
	int k;

	triplet(int x, int y, int z) {i=x; j=y; k=z;}
};

// Other Functions...
Point getVelocity(triplet ind, float *vel, triplet sz);
Point getVelocity(Point p, float *vel, triplet sz, Point lastV);
Point getVelocityInterpolated(Point p, float *vel, triplet sz);

int main(int argc, char **argv) {
	if (argc < 5) {
		std::cout << "Usage: " <<  argv[0] ; 
		// The args    1  2  3 4 5   6        7
		std::cout << " PD FA x y z sampling" << std::endl;  
		return -1;
	}

	std::cout << "Fiber Tracking Program" <<  std::endl;

	// Simple variables ...
	int x, y, z, smp;
	x = atoi(argv[3]); y = atoi(argv[4]); z = atoi(argv[5]); smp = atoi(argv[6]);
	int max_iter = 500000;
	float step_sz = 1.0;


	std::cout <<  "Allocating Memory:" << x << " " << y << " " << z; //<< tc(RST, W) << std::endl;
	// Allocate memory ...
	float * fa = new float [x*y*z];
	float * pd = new float [3*x*y*z];
	float * newpd = new float [3*x*y*z];
	std::cout <<  " Success"  << std::endl;

	// clear the new PD
	for (int i=0; i<x*y*z; i++) {
		newpd[i] = 0.0;
	}

	// Containers ...
	std::vector<Point> seeds;
	std::vector<Point> fiber;
	std::vector<Point> norms;

	std::cout <<  "Reading in anisotropy file "  << argv[2] << std::endl;
	// Read in the FA map
	std::ifstream fa_in (argv[2], std::ios_base::binary);
	fa_in.read((char *)fa, x*y*z*sizeof(float) );
	fa_in.close();

	std::cout << "Reading in principal direction "  <<  argv[1] <<  std::endl;
	// Read in the PD data
	std::ifstream pd_in (argv[1], std::ios_base::binary);
	pd_in.read((char *)pd, 3*x*y*z*sizeof(float));
	pd_in.close();

#ifdef _DEBUG
	float _mx=-100.0, _mn=100.0;
	for (int i=0; i<x*y*z; i++){
		if (fa[i] > _mx) _mx = fa[i];
		if (fa[i] < _mn) _mn = fa[i];
	}
	std::cout << "FA Range is " << _mn << " " << _mx << std::endl;
#endif

	// Compute the set of seed points based on FA ...
	for (int k=0; k<z; k+=smp) {
		for (int j=0; j<y; j+=smp) {
			for (int i=0; i<x; i+=smp) {
				if ( (fa[k*x*y + j*x +i] > 0.1 )  ) // && (fa[k*x*y + j*x + i] < 0.5))
					seeds.push_back(Point(i, j, k));
			}
		}
	}

	std::cout <<  "Computed " << seeds.size() << " Points based on anisotropy"  << std::endl; 

	clock_t t1=clock();
	// Grow fibers at each step. Stop if FA is violated. If length is longer than
	// a certain threshold, then add it as a VTK polyline ...
	for (unsigned int p=0; p<seeds.size(); p++) {
		// std::cout << p << " of " << seeds.size() << std::endl;
		fiber.clear();
		norms.clear();
		int i;

		// track for each seed
		Point lastV(0,0,0);
		Point curr = seeds[p];
		Point last = seeds[p];
		Point next = seeds[p];
		// grow in negative direction ...
		for (i=0; i<max_iter; i++) {
			Point vel = getVelocity(curr, pd, triplet(x,y,z), lastV);
			if (!i)
				vel = -vel;
			// std::cout << tc(DIM, Y, BL) << "Velocity: " << vel << tc(RST, W, BL) << std::endl;
			next = curr + vel*step_sz;
			lastV = vel;
			fiber.push_back(curr);
			norms.push_back(vel);

			if (lastV.abs() < 0.01) {
				// std::cout << tc(RST, R, BL) << "Velocity Failure" << lastV << " " << fiber.size() << tc(RST, W, BL) << std::endl;
				break;
			}
			// Test to see if the FA is all right. Later on will add additional
			// conditions at this stage like testing for curvature etc. Lets keep it
			// simple for the time being.
			if (fa[curr.getIndex(x,y)] < 0.1) {
				//std::cout << tc(RST, R, BL) << "FA failure" << fa[curr.getIndex(x,y)] << " " << fiber.size() << tc(RST, W, BL) <<  std::endl;
				break;
			}

			// move to the new location if all is well :)
			last = curr;
			curr = next;
		}
		// reverse the list, since push_front doesn't work
		std::reverse(fiber.begin(), fiber.end());
		std::reverse(norms.begin(), norms.end());
		//norms.reserve(0);

		// track for each seed
		lastV = Point(0,0,0);
		curr = seeds[p];
		last = seeds[p];
		next = seeds[p];
		// grow in positive direction ...
		for (i=0; i<max_iter; i++) {
			Point vel = getVelocity(curr, pd, triplet(x,y,z), lastV);
			// std::cout << tc(DIM, Y, BL) << "Velocity: " << vel << tc(RST, W, BL) << std::endl;
			next = curr + vel*step_sz;
			lastV = vel;
			fiber.push_back(curr);
			norms.push_back(vel);

			if (lastV.abs() < 0.01) {
				// std::cout << tc(RST, R, BL) << "Velocity Failure" << lastV << " " << fiber.size() << tc(RST, W, BL) << std::endl;
				break;
			}
			// Test to see if the FA is all right. Later on will add additional
			// conditions at this stage like testing for curvature etc. Lets keep it
			// simple for the time being.
			if (fa[curr.getIndex(x,y)] < 0.1) {
				//std::cout << tc(RST, R, BL) << "FA failure" << fa[curr.getIndex(x,y)] << " " << fiber.size() << tc(RST, W, BL) <<  std::endl;
				break;
			}

			// move to the new location if all is well :)
			last = curr;
			curr = next;
		}

		// std::cout << "For seed " << p << " got fiber of length " << fiber.size() << std::endl;
		// Now write into the newPD ...
		if (fiber.size() > 30) {
			for (unsigned int i=0; i<fiber.size(); i++) {
				int nx,ny,nz;
				nx = static_cast<int>(fiber[i].x());
				ny = static_cast<int>(fiber[i].y());
				nz = static_cast<int>(fiber[i].z());
				newpd[3*((nz*y+ny)*x + nx)] = norms[i].x();
				newpd[3*((nz*y+ny)*x + nx)+1] = norms[i].y();
				newpd[3*((nz*y+ny)*x + nx)+2] = norms[i].z();
			}
		}
	}
	clock_t t2=clock();
	std::cout << "Total time " << double(t2-t1)/CLOCKS_PER_SEC << " seconds" << std::endl;

	// write out ...
	std::ofstream pd_out ("fibers.pd", std::ios_base::binary);
	pd_out.write((char *)newpd, 3*x*y*z*sizeof(float));
	pd_out.close();

	// Clean Up ...
	std::cout << "Cleaning Up ... " <<  std::endl;
	delete [] pd;
	delete [] fa;
	delete [] newpd;

	return 0;
}


Point getVelocity (triplet ind, float *vel, triplet sz) {
	// std::cout << "getVel(int) " << ind.i << " " << ind.j << " " << ind.k << std::endl; 
	int index = 3*(ind.k*sz.i*sz.j + ind.j*sz.i + ind.i);
	return Point(vel[index], vel[index+1], vel[index+2]);
}

// Susumo's version of getVelocity at a point.
Point getVelocity (Point p, float *vel, triplet sz, Point lastV) {
	// Here we will always use the nearest neighbour to get the principal
	// direction. The vel will be scaled to hit the next boundary, and also we
	// need to test if the eigenvector needs to be inverted. This probably means
	// that we will need to send in the previous velocity estimate.

	int i, j, k;
	i = (int)p.x(); j = (int)p.y(); k = (int)p.z();

	// get PD from the floored value of p
	Point pd = getVelocity( triplet(i,j,k), vel, sz);

	// Handle the case for the first step. Need to scale the PD before returning
	// it.
	if (lastV.abs() < 0.001) {
		pd.normalize();
		return pd;
	}

	// Compare it with the last estimate of velocity. Enforce Curvature here. We
	// can change the test in the main loop to one that checks the velocity.
	// Invert the eigenvector if necessary.
	double curv = pd.dot(lastV)/(pd.abs()*lastV.abs());
	// first check if curvature is ok. If the abs is less than 0.8 then we should
	// set the velocity to zero and return. CHANGE 0.8 to a define or a static
	// variable that we read in from some config file.
	if ( fabs(curv) < 0.8 )
		return Point(0, 0, 0);

	// Since the vectors are reasonable aligned, lets flip the PD if necessary.
	if (curv < 0 ) {
		// The vectors are misaligned. Flip the PD.
		pd = - pd; // Add negation operator to the Point class.
	}

	// Scale the eigenvector to reach the next boundary. Simple first run can be
	// to simply normalize it to a unit vector which along with a step size of 1
	// should produce similar results.
	pd.normalize();

	return pd;
}

Point getVelocityInterpolated (Point p, float *vel, triplet sz) {
	// Function takes a float index and returns velocity at that point after
	// interpolation. Currently, linear interpolation is used. Might need to
	// change that eventually.
	int i, j, k;
	double x, y, z;

	i = (int)p.x(); j = (int)p.y(); k = (int)p.z();
	x = p.x() - i;  y = p.y() - j;  z = p.z() - k;

	// Simple tri-linear Interpolation ....
	Point val = getVelocity( triplet(i,j,k), vel, sz )*( (1-x)*(1-y)*(1-z) ) +
		getVelocity( triplet(i+1,j,k), vel, sz )*( x*(1-y)*(1-z) ) +
		getVelocity( triplet(i,j+1,k), vel, sz )*( (1-x)*y*(1-z) ) +
		getVelocity( triplet(i,j,k+1), vel, sz )*( (1-x)*(1-y)*z ) +
		getVelocity( triplet(i+1,j,k+1), vel, sz )*( x*(1-y)*z ) +
		getVelocity( triplet(i,j+1,k+1), vel, sz )*( (1-x)*y*z ) +
		getVelocity( triplet(i+1,j+1,k), vel, sz )*( x*y*(1-z) ) +
		getVelocity( triplet(i+1,j+1,k+1), vel, sz )*( x*y*z ) ;

	return val;
}
